#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
基于流量监控的智能重启系统
实时监控网络流量使用情况，智能决策重启时机以充分利用免费额度
"""

import time
import json
import logging
import requests
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from enum import Enum
import threading
import statistics

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RestartReason(Enum):
    """重启原因"""
    TRAFFIC_LIMIT = "traffic_limit"
    DISK_LIMIT = "disk_limit"
    SCHEDULED = "scheduled"
    EMERGENCY = "emergency"
    OPTIMIZATION = "optimization"

class RestartStatus(Enum):
    """重启状态"""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"

@dataclass
class TrafficMetrics:
    """流量指标"""
    timestamp: datetime
    monthly_usage_gb: float
    daily_usage_gb: float
    hourly_usage_gb: float
    current_rate_gb_per_hour: float
    projected_monthly_usage_gb: float
    days_remaining_in_month: int
    free_quota_remaining_gb: float

@dataclass
class InstanceInfo:
    """实例信息"""
    instance_id: str
    start_time: datetime
    disk_usage_gb: float
    network_usage_gb: float
    cpu_utilization: float
    memory_utilization: float
    status: str
    cost_accumulated: float
    active_connections: int

@dataclass
class RestartPlan:
    """重启计划"""
    instance_id: str
    reason: RestartReason
    scheduled_time: datetime
    estimated_duration: int  # 秒
    status: RestartStatus
    priority: int  # 1-5，5为最高优先级
    created_time: datetime
    execution_time: Optional[datetime] = None
    completion_time: Optional[datetime] = None
    error_message: Optional[str] = None

class TrafficMonitoringRestartSystem:
    """基于流量监控的智能重启系统"""
    
    def __init__(self):
        self.config = self._load_config()
        self.traffic_monitor = TrafficMonitor()
        self.restart_manager = RestartManager()
        self.instance_manager = InstanceManager()
        self.is_running = False
        self.monitoring_thread = None
        
    def _load_config(self) -> Dict:
        """加载配置"""
        return {
            'traffic_monitoring': {
                'check_interval': 300,  # 5分钟检查一次
                'free_quota_gb': 500,   # 免费流量额度
                'thresholds': {
                    'warning': 400,     # 400GB时警告
                    'critical': 450,    # 450GB时准备重启
                    'emergency': 480,   # 480GB时立即重启
                },
                'projection_window_hours': 24,  # 24小时预测窗口
            },
            'restart_strategy': {
                'graceful_shutdown': True,
                'drain_timeout': 60,    # 60秒排空连接
                'restart_timeout': 300, # 5分钟重启超时
                'prefer_low_traffic_hours': [2, 3, 4, 5],  # 凌晨2-5点
                'avoid_peak_hours': [9, 10, 11, 14, 15, 16, 19, 20, 21],  # 避免高峰时段
                'min_uptime_hours': 1,  # 最少运行1小时
            },
            'disk_monitoring': {
                'free_quota_gb': 80,    # 免费硬盘额度
                'thresholds': {
                    'warning': 70,      # 70GB时警告
                    'critical': 75,     # 75GB时准备重启
                },
            },
            'scheduling': {
                'max_restarts_per_day': 3,  # 每天最多重启3次
                'restart_cooldown_hours': 2,  # 重启冷却期2小时
                'maintenance_window': (2, 6),  # 维护时间窗口2-6点
            }
        }
    
    def start_monitoring(self):
        """启动监控"""
        self.is_running = True
        self.monitoring_thread = threading.Thread(target=self._monitoring_worker)
        self.monitoring_thread.daemon = True
        self.monitoring_thread.start()
        
        logger.info("流量监控重启系统已启动")
    
    def stop_monitoring(self):
        """停止监控"""
        self.is_running = False
        if self.monitoring_thread:
            self.monitoring_thread.join()
        
        logger.info("流量监控重启系统已停止")
    
    def _monitoring_worker(self):
        """监控工作线程"""
        while self.is_running:
            try:
                # 更新流量指标
                traffic_metrics = self.traffic_monitor.get_current_metrics()
                
                # 分析是否需要重启
                restart_decision = self._analyze_restart_need(traffic_metrics)
                
                if restart_decision['should_restart']:
                    # 创建重启计划
                    restart_plan = self._create_restart_plan(restart_decision)
                    
                    # 执行重启计划
                    self._execute_restart_plan(restart_plan)
                
                # 等待下次检查
                time.sleep(self.config['traffic_monitoring']['check_interval'])
                
            except Exception as e:
                logger.error(f"监控工作线程异常: {e}")
                time.sleep(60)
    
    def _analyze_restart_need(self, traffic_metrics: TrafficMetrics) -> Dict:
        """分析是否需要重启"""
        decision = {
            'should_restart': False,
            'reason': None,
            'urgency': 'low',
            'scheduled_time': None,
            'priority': 1,
        }
        
        # 检查流量使用情况
        if traffic_metrics.monthly_usage_gb >= self.config['traffic_monitoring']['thresholds']['emergency']:
            decision.update({
                'should_restart': True,
                'reason': RestartReason.TRAFFIC_LIMIT,
                'urgency': 'emergency',
                'priority': 5,
            })
        elif traffic_metrics.monthly_usage_gb >= self.config['traffic_monitoring']['thresholds']['critical']:
            decision.update({
                'should_restart': True,
                'reason': RestartReason.TRAFFIC_LIMIT,
                'urgency': 'high',
                'priority': 4,
            })
        elif traffic_metrics.projected_monthly_usage_gb > self.config['traffic_monitoring']['free_quota_gb']:
            decision.update({
                'should_restart': True,
                'reason': RestartReason.TRAFFIC_LIMIT,
                'urgency': 'medium',
                'priority': 3,
            })
        
        # 检查硬盘使用情况
        instance_info = self.instance_manager.get_current_instance_info()
        if instance_info and instance_info.disk_usage_gb >= self.config['disk_monitoring']['thresholds']['critical']:
            decision.update({
                'should_restart': True,
                'reason': RestartReason.DISK_LIMIT,
                'urgency': 'high',
                'priority': 4,
            })
        
        # 检查运行时间
        if instance_info:
            uptime_hours = (datetime.now() - instance_info.start_time).total_seconds() / 3600
            if uptime_hours >= 24 * 15:  # 运行超过15天
                decision.update({
                    'should_restart': True,
                    'reason': RestartReason.SCHEDULED,
                    'urgency': 'low',
                    'priority': 2,
                })
        
        # 确定重启时间
        if decision['should_restart']:
            decision['scheduled_time'] = self._choose_optimal_restart_time(decision['urgency'])
        
        return decision
    
    def _choose_optimal_restart_time(self, urgency: str) -> datetime:
        """选择最佳重启时间"""
        current_time = datetime.now()
        current_hour = current_time.hour
        
        if urgency == 'emergency':
            # 紧急情况，立即重启
            return current_time + timedelta(minutes=5)
        
        # 优先选择低流量时段
        low_traffic_hours = self.config['restart_strategy']['prefer_low_traffic_hours']
        avoid_peak_hours = self.config['restart_strategy']['avoid_peak_hours']
        
        # 检查当前是否在低流量时段
        if current_hour in low_traffic_hours and current_hour not in avoid_peak_hours:
            return current_time + timedelta(minutes=10)
        
        # 计算下一个低流量时段
        next_low_traffic_hour = None
        for hour in low_traffic_hours:
            if hour > current_hour and hour not in avoid_peak_hours:
                next_low_traffic_hour = hour
                break
        
        if next_low_traffic_hour is None:
            # 如果今天没有合适的时段，选择明天最早的低流量时段
            next_low_traffic_hour = min(low_traffic_hours)
            target_time = current_time.replace(hour=next_low_traffic_hour, minute=0, second=0) + timedelta(days=1)
        else:
            target_time = current_time.replace(hour=next_low_traffic_hour, minute=0, second=0)
        
        return target_time
    
    def _create_restart_plan(self, decision: Dict) -> RestartPlan:
        """创建重启计划"""
        instance_info = self.instance_manager.get_current_instance_info()
        
        plan = RestartPlan(
            instance_id=instance_info.instance_id if instance_info else "unknown",
            reason=decision['reason'],
            scheduled_time=decision['scheduled_time'],
            estimated_duration=300,  # 5分钟
            status=RestartStatus.PENDING,
            priority=decision['priority'],
            created_time=datetime.now()
        )
        
        logger.info(f"创建重启计划: {plan.instance_id}, 原因: {plan.reason.value}, 时间: {plan.scheduled_time}")
        
        return plan
    
    def _execute_restart_plan(self, plan: RestartPlan):
        """执行重启计划"""
        try:
            # 检查是否可以执行重启
            if not self._can_execute_restart(plan):
                logger.info(f"重启计划被取消: {plan.instance_id}")
                return
            
            # 更新计划状态
            plan.status = RestartStatus.IN_PROGRESS
            plan.execution_time = datetime.now()
            
            logger.info(f"开始执行重启计划: {plan.instance_id}")
            
            # 执行重启流程
            success = self._perform_restart(plan)
            
            if success:
                plan.status = RestartStatus.COMPLETED
                plan.completion_time = datetime.now()
                logger.info(f"重启计划执行成功: {plan.instance_id}")
            else:
                plan.status = RestartStatus.FAILED
                plan.completion_time = datetime.now()
                logger.error(f"重启计划执行失败: {plan.instance_id}")
                
        except Exception as e:
            plan.status = RestartStatus.FAILED
            plan.error_message = str(e)
            plan.completion_time = datetime.now()
            logger.error(f"执行重启计划异常: {plan.instance_id}, {e}")
    
    def _can_execute_restart(self, plan: RestartPlan) -> bool:
        """检查是否可以执行重启"""
        # 检查冷却期
        last_restart_time = self.restart_manager.get_last_restart_time()
        if last_restart_time:
            cooldown_hours = self.config['scheduling']['restart_cooldown_hours']
            time_since_last_restart = (datetime.now() - last_restart_time).total_seconds() / 3600
            
            if time_since_last_restart < cooldown_hours:
                logger.info(f"重启冷却期中，距离上次重启 {time_since_last_restart:.1f} 小时")
                return False
        
        # 检查每日重启次数限制
        today_restart_count = self.restart_manager.get_today_restart_count()
        max_restarts = self.config['scheduling']['max_restarts_per_day']
        
        if today_restart_count >= max_restarts:
            logger.info(f"今日重启次数已达上限: {today_restart_count}/{max_restarts}")
            return False
        
        # 检查最小运行时间
        instance_info = self.instance_manager.get_current_instance_info()
        if instance_info:
            uptime_hours = (datetime.now() - instance_info.start_time).total_seconds() / 3600
            min_uptime = self.config['restart_strategy']['min_uptime_hours']
            
            if uptime_hours < min_uptime:
                logger.info(f"实例运行时间不足: {uptime_hours:.1f}小时 < {min_uptime}小时")
                return False
        
        return True
    
    def _perform_restart(self, plan: RestartPlan) -> bool:
        """执行重启"""
        try:
            instance_info = self.instance_manager.get_current_instance_info()
            if not instance_info:
                logger.error("无法获取实例信息")
                return False
            
            # 1. 优雅关闭
            if self.config['restart_strategy']['graceful_shutdown']:
                logger.info(f"开始优雅关闭实例: {plan.instance_id}")
                success = self.instance_manager.graceful_shutdown(plan.instance_id)
                if not success:
                    logger.error("优雅关闭失败")
                    return False
            
            # 2. 等待连接排空
            drain_timeout = self.config['restart_strategy']['drain_timeout']
            logger.info(f"等待连接排空: {drain_timeout}秒")
            time.sleep(drain_timeout)
            
            # 3. 创建新实例
            logger.info("创建新实例")
            new_instance = self.instance_manager.create_new_instance()
            if not new_instance:
                logger.error("创建新实例失败")
                return False
            
            # 4. 验证新实例
            logger.info(f"验证新实例: {new_instance['instance_id']}")
            if not self.instance_manager.verify_instance(new_instance['instance_id']):
                logger.error("新实例验证失败")
                return False
            
            # 5. 销毁旧实例
            logger.info(f"销毁旧实例: {plan.instance_id}")
            if not self.instance_manager.destroy_instance(plan.instance_id):
                logger.error("销毁旧实例失败")
                return False
            
            # 6. 记录重启成功
            self.restart_manager.record_restart(plan)
            
            return True
            
        except Exception as e:
            logger.error(f"执行重启异常: {e}")
            return False
    
    def get_system_status(self) -> Dict:
        """获取系统状态"""
        traffic_metrics = self.traffic_monitor.get_current_metrics()
        instance_info = self.instance_manager.get_current_instance_info()
        
        return {
            'traffic_metrics': {
                'monthly_usage_gb': traffic_metrics.monthly_usage_gb,
                'projected_monthly_usage_gb': traffic_metrics.projected_monthly_usage_gb,
                'free_quota_remaining_gb': traffic_metrics.free_quota_remaining_gb,
                'days_remaining_in_month': traffic_metrics.days_remaining_in_month,
            },
            'instance_info': {
                'instance_id': instance_info.instance_id if instance_info else None,
                'uptime_hours': (datetime.now() - instance_info.start_time).total_seconds() / 3600 if instance_info else 0,
                'disk_usage_gb': instance_info.disk_usage_gb if instance_info else 0,
                'network_usage_gb': instance_info.network_usage_gb if instance_info else 0,
            },
            'restart_history': self.restart_manager.get_recent_restarts(10),
            'next_restart_estimate': self._estimate_next_restart_time(),
        }
    
    def _estimate_next_restart_time(self) -> Optional[datetime]:
        """估算下次重启时间"""
        traffic_metrics = self.traffic_monitor.get_current_metrics()
        
        if traffic_metrics.projected_monthly_usage_gb > self.config['traffic_monitoring']['free_quota_gb']:
            # 基于流量使用速率计算
            if traffic_metrics.current_rate_gb_per_hour > 0:
                hours_to_limit = (self.config['traffic_monitoring']['free_quota_gb'] - traffic_metrics.monthly_usage_gb) / traffic_metrics.current_rate_gb_per_hour
                return datetime.now() + timedelta(hours=hours_to_limit)
        
        return None

class TrafficMonitor:
    """流量监控器"""
    
    def __init__(self):
        self.metrics_history = []
    
    def get_current_metrics(self) -> TrafficMetrics:
        """获取当前流量指标"""
        # 这里应该从实际的监控系统获取数据
        # 模拟数据
        current_time = datetime.now()
        monthly_usage = 420.5  # 模拟月使用量
        daily_usage = 15.2     # 模拟日使用量
        hourly_usage = 0.8     # 模拟小时使用量
        
        # 计算使用速率
        current_rate = hourly_usage
        
        # 计算剩余天数
        days_remaining = 30 - current_time.day
        
        # 预测月使用量
        projected_monthly = monthly_usage + (current_rate * 24 * days_remaining)
        
        # 计算剩余免费额度
        free_quota_remaining = 500 - monthly_usage
        
        metrics = TrafficMetrics(
            timestamp=current_time,
            monthly_usage_gb=monthly_usage,
            daily_usage_gb=daily_usage,
            hourly_usage_gb=hourly_usage,
            current_rate_gb_per_hour=current_rate,
            projected_monthly_usage_gb=projected_monthly,
            days_remaining_in_month=days_remaining,
            free_quota_remaining_gb=free_quota_remaining
        )
        
        # 保存历史记录
        self.metrics_history.append(metrics)
        if len(self.metrics_history) > 1000:
            self.metrics_history = self.metrics_history[-500:]
        
        return metrics

class RestartManager:
    """重启管理器"""
    
    def __init__(self):
        self.restart_history = []
    
    def record_restart(self, plan: RestartPlan):
        """记录重启"""
        self.restart_history.append(plan)
        logger.info(f"记录重启: {plan.instance_id}, 原因: {plan.reason.value}")
    
    def get_last_restart_time(self) -> Optional[datetime]:
        """获取上次重启时间"""
        if not self.restart_history:
            return None
        
        last_restart = max(self.restart_history, key=lambda x: x.completion_time or x.execution_time or x.created_time)
        return last_restart.completion_time or last_restart.execution_time or last_restart.created_time
    
    def get_today_restart_count(self) -> int:
        """获取今日重启次数"""
        today = datetime.now().date()
        count = 0
        
        for restart in self.restart_history:
            restart_date = (restart.completion_time or restart.execution_time or restart.created_time).date()
            if restart_date == today:
                count += 1
        
        return count
    
    def get_recent_restarts(self, limit: int = 10) -> List[Dict]:
        """获取最近的重启记录"""
        recent_restarts = sorted(self.restart_history, 
                               key=lambda x: x.completion_time or x.execution_time or x.created_time, 
                               reverse=True)[:limit]
        
        return [
            {
                'instance_id': r.instance_id,
                'reason': r.reason.value,
                'status': r.status.value,
                'created_time': r.created_time.isoformat(),
                'execution_time': r.execution_time.isoformat() if r.execution_time else None,
                'completion_time': r.completion_time.isoformat() if r.completion_time else None,
                'error_message': r.error_message,
            }
            for r in recent_restarts
        ]

class InstanceManager:
    """实例管理器"""
    
    def __init__(self):
        self.current_instance = None
    
    def get_current_instance_info(self) -> Optional[InstanceInfo]:
        """获取当前实例信息"""
        # 这里应该从实际的实例管理API获取信息
        # 模拟数据
        if not self.current_instance:
            self.current_instance = InstanceInfo(
                instance_id="hai-instance-001",
                start_time=datetime.now() - timedelta(days=10),
                disk_usage_gb=65.5,
                network_usage_gb=420.5,
                cpu_utilization=45.0,
                memory_utilization=60.0,
                status="RUNNING",
                cost_accumulated=25.6,
                active_connections=5
            )
        
        return self.current_instance
    
    def graceful_shutdown(self, instance_id: str) -> bool:
        """优雅关闭实例"""
        logger.info(f"优雅关闭实例: {instance_id}")
        # 这里应该调用实际的API
        time.sleep(2)  # 模拟关闭过程
        return True
    
    def create_new_instance(self) -> Optional[Dict]:
        """创建新实例"""
        logger.info("创建新实例")
        # 这里应该调用实际的API
        time.sleep(5)  # 模拟创建过程
        
        new_instance = {
            'instance_id': f"hai-instance-{int(time.time())}",
            'status': 'RUNNING',
            'created_time': datetime.now()
        }
        
        return new_instance
    
    def verify_instance(self, instance_id: str) -> bool:
        """验证实例"""
        logger.info(f"验证实例: {instance_id}")
        # 这里应该调用实际的API验证实例状态
        time.sleep(2)  # 模拟验证过程
        return True
    
    def destroy_instance(self, instance_id: str) -> bool:
        """销毁实例"""
        logger.info(f"销毁实例: {instance_id}")
        # 这里应该调用实际的API
        time.sleep(3)  # 模拟销毁过程
        return True

# 使用示例
def main():
    """使用示例"""
    system = TrafficMonitoringRestartSystem()
    
    try:
        # 启动监控
        system.start_monitoring()
        
        # 运行一段时间
        time.sleep(60)
        
        # 获取系统状态
        status = system.get_system_status()
        print(f"系统状态: {json.dumps(status, indent=2, default=str)}")
        
    finally:
        # 停止监控
        system.stop_monitoring()

if __name__ == '__main__':
    main()
